import tkinter as tk
from unicodedata import digit
import numpy as np

def draw_digit():
    """
    Opens a window for the user to draw a digit like in the MNIST dataset.
    Returns a 784-element numpy vector with values between 0 and 1 when the user clicks "Save".
    """
    # Create a result variable to store the final vector
    result = {"vector": None}
    
    # Create the root window
    root = tk.Tk()
    root.title("MNIST Digit Drawing")
    root.resizable(False, False)
    
    # Constants
    MNIST_DIM = 28
    SCALE = 10  # Scale up for better drawing experience
    CANVAS_DIM = MNIST_DIM * SCALE
    
    # Store the pixel values (0-1 float, grayscale)
    pixels = np.zeros((MNIST_DIM, MNIST_DIM))
    
    def update_status(event):
        x, y = event.x // SCALE, event.y // SCALE
        if 0 <= x < MNIST_DIM and 0 <= y < MNIST_DIM:
            pixel_value = pixels[y, x]
            status_var.set(f"Position: ({x}, {y}) - Value: {pixel_value:.2f}")
    
    def draw_grid():
        # Draw grid lines
        for i in range(MNIST_DIM + 1):
            # Vertical lines
            x = i * SCALE
            canvas.create_line(x, 0, x, CANVAS_DIM, fill="#333333")
            
            # Horizontal lines
            y = i * SCALE
            canvas.create_line(0, y, CANVAS_DIM, y, fill="#333333")
    
    def paint(event):
        # Get the pen size and intensity
        pen_size = pen_size_var.get()
        intensity = intensity_var.get()
        
        # Calculate center pixel in our grid coordinates
        center_x = event.x // SCALE
        center_y = event.y // SCALE
        
        # Apply the pen with a circular brush and falloff
        for y_offset in range(-pen_size, pen_size + 1):
            for x_offset in range(-pen_size, pen_size + 1):
                x = center_x + x_offset
                y = center_y + y_offset
                
                # Skip if outside the grid
                if not (0 <= x < MNIST_DIM and 0 <= y < MNIST_DIM):
                    continue
                
                # Calculate distance from center (for circular brush)
                distance = (x_offset**2 + y_offset**2) ** 0.5
                
                # Skip if outside the circular brush
                if distance > pen_size:
                    continue
                
                # Calculate falloff based on distance from center
                falloff = 1 - (distance / pen_size)
                
                # Update pixel value with current intensity, with falloff
                current_value = pixels[y, x]
                new_value = min(1.0, current_value + intensity * falloff)
                pixels[y, x] = new_value
                
                # Update the visual representation
                fill_color = _get_grayscale_color(new_value)
                
                canvas.create_rectangle(
                    x * SCALE, y * SCALE,
                    (x + 1) * SCALE, (y + 1) * SCALE,
                    fill=fill_color, outline=""
                )
        
        # Update the preview
        update_preview()
        
        # Update status bar
        update_status(event)
    
    def _get_grayscale_color(value):
        # Convert float value (0-1) to hex color code
        color_value = int(value * 255)
        return f'#{color_value:02x}{color_value:02x}{color_value:02x}'
    
    def clear_canvas():
        # Clear the canvas and reset pixel values
        canvas.delete("all")
        pixels.fill(0)
        draw_grid()
        update_preview()
        status_var.set("Canvas cleared")
    
    def update_preview():
        # Clear the preview canvas
        preview_canvas.delete("all")
        
        # Draw the preview (scaled up 4x, inverted colors for MNIST-like view)
        preview_scale = 4
        for y in range(MNIST_DIM):
            for x in range(MNIST_DIM):
                # If the pixel has a value, draw it (inverted for MNIST-like look)
                if pixels[y, x] > 0:
                    # Invert colors for the preview (black on white like MNIST)
                    color_value = int((1 - pixels[y, x]) * 255)
                    color = f'#{color_value:02x}{color_value:02x}{color_value:02x}'
                    
                    preview_canvas.create_rectangle(
                        x * preview_scale, y * preview_scale, 
                        (x + 1) * preview_scale, (y + 1) * preview_scale,
                        fill=color, outline=""
                    )
    
    def save_and_exit():
        # Store the flattened pixel array in the result
        result["vector"] = pixels.flatten()
        # Close the window
        root.destroy()
    
    # Main frame
    main_frame = tk.Frame(root)
    main_frame.pack(padx=10, pady=10)
    
    # Left panel (controls)
    left_panel = tk.Frame(main_frame)
    left_panel.pack(side=tk.LEFT, padx=10)
    
    # Drawing controls
    tk.Label(left_panel, text="Drawing Controls").pack(pady=(0, 5))
    
    # Intensity selection
    intensity_frame = tk.Frame(left_panel)
    intensity_frame.pack(pady=5, fill=tk.X)
    
    tk.Label(intensity_frame, text="Intensity:").pack(side=tk.LEFT)
    intensity_var = tk.DoubleVar(value=1.0)
    intensity_slider = tk.Scale(intensity_frame, from_=0.1, to=1.0, 
                              resolution=0.1, orient=tk.HORIZONTAL,
                              variable=intensity_var)
    intensity_slider.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
    
    # Pen size selection
    pen_frame = tk.Frame(left_panel)
    pen_frame.pack(pady=5, fill=tk.X)
    
    tk.Label(pen_frame, text="Pen Size:").pack(side=tk.LEFT)
    pen_size_var = tk.IntVar(value=2)
    pen_slider = tk.Scale(pen_frame, from_=1, to=5, 
                         orient=tk.HORIZONTAL,
                         variable=pen_size_var)
    pen_slider.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
    
    # Buttons frame - define it BEFORE using it
    button_frame = tk.Frame(left_panel)
    button_frame.pack(pady=10, fill=tk.X)
    
    # Add the clear and save buttons
    clear_btn = tk.Button(button_frame, text="Clear", command=clear_canvas)
    clear_btn.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
    
    save_btn = tk.Button(button_frame, text="Save", command=save_and_exit)
    save_btn.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
    
    # Right panel (drawing area)
    right_panel = tk.Frame(main_frame)
    right_panel.pack(side=tk.LEFT, padx=10)
    
    # Canvas setup
    canvas_frame = tk.Frame(right_panel, bd=2, relief=tk.SUNKEN)
    canvas_frame.pack(pady=5)
    
    canvas = tk.Canvas(canvas_frame, width=CANVAS_DIM, height=CANVAS_DIM, 
                     bg="black", highlightthickness=0)
    canvas.pack()
    
    # Preview
    preview_frame = tk.LabelFrame(right_panel, text="MNIST Preview (28x28)")
    preview_frame.pack(pady=10, fill=tk.BOTH, expand=True)
    
    preview_canvas = tk.Canvas(preview_frame, width=MNIST_DIM*4, 
                             height=MNIST_DIM*4, bg="white")
    preview_canvas.pack(padx=10, pady=10)
    
    # Status bar
    status_frame = tk.Frame(root, bd=1, relief=tk.SUNKEN)
    status_frame.pack(side=tk.BOTTOM, fill=tk.X)
    
    status_var = tk.StringVar(value="Ready to draw")
    status_label = tk.Label(status_frame, textvariable=status_var, anchor=tk.W)
    status_label.pack(fill=tk.X, padx=5, pady=2)
    
    # Bind mouse events
    canvas.bind("<B1-Motion>", paint)
    canvas.bind("<ButtonPress-1>", paint)
    canvas.bind("<Motion>", update_status)
    
    # Draw the grid initially
    draw_grid()
    
    # Start the main loop
    root.mainloop()
    
    # Return the result after the window is closed
    return result["vector"]

# Example of how to use this function in another program:
if __name__ == "__main__":
    # This is how you call the function
    digit_vector = draw_digit()
    
    # Once the user clicks "Save", the function returns the vector
    if digit_vector is not None:
        print("Received vector with shape:", digit_vector.shape)
        print("First 10 values:", digit_vector[:10])
        digit_vector = digit_vector.reshape(28, -1)
        np.set_printoptions(linewidth=150)
        print(np.round(digit_vector))
        
        # You can use this vector in your neural network, classifier, etc.
        # For example:
        # prediction = your_model.predict(digit_vector.reshape(1, 784))